//
// Created by cybaster on 25-5-27.
//

#include <angles/angles.h>
#include <cstring>

#include <ocs2_centroidal_model/AccessHelperFunctions.h>
#include <ocs2_core/thread_support/ExecuteAndSleep.h>
#include <ocs2_core/thread_support/SetThreadPriority.h>
#include <ocs2_robotic_tools/common/RotationDerivativesTransforms.h>
#include <ocs2_ros_interfaces/common/RosMsgConversions.h>
#include <ocs2_ros_interfaces/synchronized_module/RosReferenceManager.h>
#include <ocs2_sqp/SqpMpc.h>

#include <std_msgs/Float64.h>
#include <std_msgs/Float64MultiArray.h>

#include "mpc_quadruped_mujoco/controller/LeggedRobotController.h"
#include "mpc_quadruped_mujoco/gait/GaitReceiver.h"

double sign(scalar_t x) {
  if (x > 0.0) {
    return 1.0;
  }
  if (x < 0.0) {
    return -1.0;
  }
  return 0.0;
}

vector3_t wrapEulerAnglesZyx(const vector3_t &eulerAngles) {
  vector3_t wrappedEulerAngles = eulerAngles;
  if (fabs(wrappedEulerAngles(1)) > M_PI_2) {
    wrappedEulerAngles(1) =
        sign(wrappedEulerAngles(1)) * M_PI - wrappedEulerAngles(1);
    wrappedEulerAngles(0) =
        wrappedEulerAngles(0) - sign(wrappedEulerAngles(0)) * M_PI;
    wrappedEulerAngles(2) =
        wrappedEulerAngles(2) - sign(wrappedEulerAngles(2)) * M_PI;
  }
  return wrappedEulerAngles;
}

LeggedRobotController::LeggedRobotController(
    LeggedRobotInterface &leggedRobotInterface)
    : interface_(&leggedRobotInterface),
      info_(leggedRobotInterface.getCentroidalModelInfo()) {

  robotName_ = "legged_robot";
  timestep = 1.0 / interface_->mpcSettings().mrtDesiredFrequency_;

  jointPos_ = interface_->getInitialState().tail(info_.actuatedDofNum);
  jointVel_ = vector_t::Zero(info_.actuatedDofNum);

  pos_ = interface_->getInitialState().segment(6, 3);
  vector3_t eulerAnglesZyx;
  eulerAnglesZyx = interface_->getInitialState().segment(9, 3);
  quat_ =
      Eigen::AngleAxis<scalar_t>(eulerAnglesZyx(0), Eigen::Vector3d::UnitZ()) *
      Eigen::AngleAxis<scalar_t>(eulerAnglesZyx(1), Eigen::Vector3d::UnitY()) *
      Eigen::AngleAxis<scalar_t>(eulerAnglesZyx(2), Eigen::Vector3d::UnitX());
  angularVelLocal_.setZero();
  linearVelLocal_.setZero();

  rbdConversions_ = std::make_shared<CentroidalModelRbdConversions>(
      interface_->getPinocchioInterface(),
      interface_->getCentroidalModelInfo());
  rbdState_ = vector_t::Zero(2 * info_.generalizedCoordinatesNum);

  setupMpc();
  setupMrt();
}

void LeggedRobotController::launchNodes(ros::NodeHandle &nodeHandle) {
  auto gaitReceiverPtr = std::make_shared<GaitReceiver>(
      nodeHandle,
      interface_->getSwitchedModelReferenceManagerPtr()->getGaitSchedule(),
      robotName_);
  auto rosReferenceManagerPtr = std::make_shared<RosReferenceManager>(
      robotName_, interface_->getReferenceManagerPtr());
  rosReferenceManagerPtr->subscribe(nodeHandle);
  mpc_->getSolverPtr()->addSynchronizedModule(gaitReceiverPtr);
  mpc_->getSolverPtr()->setReferenceManager(rosReferenceManagerPtr);

  imuSubscriber_ = nodeHandle.subscribe<sensor_msgs::Imu>(
      robotName_ + "_imu", 1, &LeggedRobotController::imuCallback, this);
  odometrySubscriber_ = nodeHandle.subscribe<nav_msgs::Odometry>(
      robotName_ + "_odometry", 1, &LeggedRobotController::odometryCallback,
      this);
  jointStateSubscriber_ = nodeHandle.subscribe<sensor_msgs::JointState>(
      robotName_ + "_joint_state", 1,
      &LeggedRobotController::jointStateCallback, this);
  mpcObservationPublisher_ = nodeHandle.advertise<ocs2_msgs::mpc_observation>(
      robotName_ + "_mpc_observation", 1);
  mpcCtrlPublisher_ = nodeHandle.advertise<std_msgs::Float64MultiArray>(
      robotName_ + "_mpc_ctrl", 1);

  resetServer_ =
      nodeHandle.advertiseService(robotName_ + "_mpc_reset_mpc",
                                  &LeggedRobotController::resetCallback, this);

  ros::Rate rate(interface_->mpcSettings().mrtDesiredFrequency_);
  timer_ = nodeHandle.createTimer(ros::Duration(rate),
                                  &LeggedRobotController::timerCallback, this);

  ros::spin();
  ros::shutdown();
}

void LeggedRobotController::setupMpc() {
  mpc_ = std::make_shared<SqpMpc>(
      interface_->mpcSettings(), interface_->sqpSettings(),
      interface_->getOptimalControlProblem(), interface_->getInitializer());
}

void LeggedRobotController::setupMrt() {
  mpcMrtInterface_ = std::make_shared<MPC_MRT_Interface>(*mpc_);
  mpcMrtInterface_->initRollout(&interface_->getRollout());

  controllerRunning_ = true;
  mpcThread_ = std::thread([&]() {
    while (controllerRunning_) {
      try {
        executeAndSleep(
            [&]() {
              if (mpcRunning_) {
                mpcMrtInterface_->advanceMpc();
              }
            },
            interface_->mpcSettings().mpcDesiredFrequency_);
      } catch (std::exception &e) {
        controllerRunning_ = false;
        ROS_ERROR_STREAM("[Ocs2 MPC thread] Error : " << e.what());
      }
    }
  });
  setThreadPriority(interface_->sqpSettings().threadPriority, mpcThread_);
}

void LeggedRobotController::start() {
  currentObservation_.time = 0.0;
  currentObservation_.state = interface_->getInitialState();
  currentObservation_.input.setZero(info_.inputDim);
  currentObservation_.mode = STANCE;
  plannerMode_ = currentObservation_.mode;

  TargetTrajectories targetTrajectories({currentObservation_.time},
                                        {currentObservation_.state},
                                        {currentObservation_.input});
  mpcMrtInterface_->setCurrentObservation(currentObservation_);
  mpcMrtInterface_->getReferenceManager().setTargetTrajectories(
      targetTrajectories);
  ROS_INFO_STREAM("Waiting for the initial policy ...");
  while (!mpcMrtInterface_->initialPolicyReceived() && ros::ok()) {
    mpcMrtInterface_->advanceMpc();
    ros::WallRate(interface_->mpcSettings().mrtDesiredFrequency_).sleep();
  }
  ROS_INFO_STREAM("Initial policy has been received.");

  mpcRunning_ = true;
}

void LeggedRobotController::stop() { mpcRunning_ = false; }

void LeggedRobotController::update() {
  updateState();

  mpcMrtInterface_->setCurrentObservation(currentObservation_);
  mpcMrtInterface_->updatePolicy();

  vector_t optimizedState, optimizedInput;
  size_t plannerMode = 0;
  mpcMrtInterface_->evaluatePolicy(currentObservation_.time,
                                   currentObservation_.state, optimizedState,
                                   optimizedInput, plannerMode);

  currentObservation_.input = optimizedInput;

  vector_t jointAccelerations(info_.actuatedDofNum);
  jointAccelerations.setZero();
  auto jointTorque = rbdConversions_
                         ->computeRbdTorqueFromCentroidalModel(
                             optimizedState, optimizedInput, jointAccelerations)
                         .tail(info_.actuatedDofNum);

  auto jointPos = centroidal_model::getJointAngles(optimizedState, info_);
  auto jointVel = centroidal_model::getJointVelocities(optimizedInput, info_);

  std_msgs::Float64MultiArray mpcCtrl;
  for (int i = 0; i < info_.actuatedDofNum; ++i) {
    mpcCtrl.data.push_back(jointPos(i));
  }
  for (int i = 0; i < info_.actuatedDofNum; ++i) {
    mpcCtrl.data.push_back(jointVel(i));
  }
  for (int i = 0; i < info_.actuatedDofNum; ++i) {
    mpcCtrl.data.push_back(jointTorque(i));
  }

  mpcCtrlPublisher_.publish(mpcCtrl);
  mpcObservationPublisher_.publish(
      ros_msg_conversions::createObservationMsg(currentObservation_));
}

void LeggedRobotController::updateState() {

  vector3_t pos, linearVelLocal, angularVelLocal;
  Eigen::Quaternion<scalar_t> quat;

  {
    std::lock_guard<std::mutex> lock(imuMutex_);
    quat = quat_;
    angularVelLocal = angularVelLocal_;
  }
  {
    std::lock_guard<std::mutex> lock(odometryMutex_);
    pos = pos_;
    linearVelLocal = linearVelLocal_;
  }

  auto eulerAnglesZyx =
      wrapEulerAnglesZyx(quat.toRotationMatrix().eulerAngles(2, 1, 0));

  matrix3_t rotationMatrix = quat.toRotationMatrix();
  auto linearVelGlobal = rotationMatrix * linearVelLocal;

  auto angularVelGlobal = getGlobalAngularVelocityFromEulerAnglesZyxDerivatives(
      eulerAnglesZyx, getEulerAnglesZyxDerivativesFromLocalAngularVelocity(
                          eulerAnglesZyx, angularVelLocal));

  rbdState_.segment<3>(0) = eulerAnglesZyx;
  rbdState_.segment<3>(info_.generalizedCoordinatesNum) = angularVelGlobal;
  rbdState_.segment<3>(3) = pos;
  rbdState_.segment<3>(info_.generalizedCoordinatesNum + 3) = linearVelGlobal;

  {
    std::lock_guard<std::mutex> lock(jointMutex_);
    rbdState_.segment(6, info_.actuatedDofNum) = jointPos_;
    rbdState_.segment(6 + info_.generalizedCoordinatesNum,
                      info_.actuatedDofNum) = jointVel_;
  }

  currentObservation_.time += timestep;
  auto yawLast = currentObservation_.state(9);
  currentObservation_.state =
      rbdConversions_->computeCentroidalStateFromRbdModel(rbdState_);
  currentObservation_.state(9) =
      yawLast +
      angles::shortest_angular_distance(yawLast, currentObservation_.state(9));
  currentObservation_.mode = plannerMode_;
}

void LeggedRobotController::jointStateCallback(
    const sensor_msgs::JointState::ConstPtr &msg) {
  std::lock_guard<std::mutex> lock(jointMutex_);
  for (size_t i = 0; i < info_.actuatedDofNum; i++) {
    jointPos_(i) = msg->position[i];
    jointVel_(i) = msg->velocity[i];
  }
}

void LeggedRobotController::imuCallback(const sensor_msgs::Imu::ConstPtr &msg) {
  std::lock_guard<std::mutex> lock(imuMutex_);
  quat_ = Eigen::Quaternion<scalar_t>(msg->orientation.w, msg->orientation.x,
                                      msg->orientation.y, msg->orientation.z);
  angularVelLocal_ << msg->angular_velocity.x, msg->angular_velocity.y,
      msg->angular_velocity.z;
}

void LeggedRobotController::odometryCallback(
    const nav_msgs::Odometry::ConstPtr &msg) {
  std::lock_guard<std::mutex> lock(odometryMutex_);
  pos_ << msg->pose.pose.position.x, msg->pose.pose.position.y,
      msg->pose.pose.position.z;
  linearVelLocal_ << msg->twist.twist.linear.x, msg->twist.twist.linear.y,
      msg->twist.twist.linear.z;
}

void LeggedRobotController::timerCallback(const ros::TimerEvent &event) {
  if (!mpcRunning_) {
    return;
  }

  update();
}

bool LeggedRobotController::resetCallback(std_srvs::SetBool::Request &req,
                                          std_srvs::SetBool::Response &res) {
  start();
  res.success = true;
  return true;
}
